# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
# RUN: %mojo-build %s -o %t
# RUN: %t

from os import getenv, listdir, setenv
from os.path import dirname
from pathlib import Path, cwd
from subprocess import run
from sys.ffi import c_int
from sys.param_env import env_get_string

from gpu.host.device_attribute import DeviceAttribute
from gpu.host.dim import Dim
from gpu.host import DeviceBuffer
from memory import LegacyUnsafePointer as UnsafePointer
from python import Python
from shmem import *
from testing import assert_equal


fn ring_bcast(
    data: UnsafePointer[c_int],
    nelem: Int,
    root: c_int,
    psync: UnsafePointer[UInt64],
):
    var mype = shmem_my_pe()
    var npes = shmem_n_pes()
    var peer = (mype + 1) % npes

    if mype == root:
        psync[0] = 1

    shmem_signal_wait_until(psync, SHMEM_CMP_NE, 0)

    if mype == npes - 1:
        return

    shmem_put(data, data, UInt(nelem), peer)
    shmem_fence()
    shmem_signal_op(psync, 1, SHMEM_SIGNAL_SET, peer)

    psync[0] = 0


def test_ring_bcast(ctx: SHMEMContext):
    alias data_len = 32
    var destination = ctx.enqueue_create_buffer[DType.int32](1)

    var data = ctx.enqueue_create_buffer[DType.int32](data_len)
    var data_h = UnsafePointer[Int32].alloc(data_len)
    var psync = shmem_calloc[DType.uint64](1)

    for i in range(data_len):
        data_h[i] = shmem_my_pe() + i

    data.enqueue_copy_from(data_h)

    var root: Int32 = 0
    ctx.barrier_all()
    ctx.enqueue_function_collective_checked[ring_bcast, ring_bcast](
        data,
        data_len,
        root,
        DeviceBuffer[DType.uint64](ctx._ctx, psync, 1, owning=False),
        grid_dim=1,
        block_dim=1,
    )
    ctx.barrier_all()

    data.enqueue_copy_to(data_h)
    ctx.synchronize()

    var mype = shmem_my_pe()
    for i in range(data_len):
        assert_equal(
            data_h[i],
            Int32(i),
            String(
                "PE",
                mype,
                "error, data[",
                i,
                "] = ",
                data_h[i],
                " expected ",
                i,
            ),
        )


def main():
    shmem_launch[test_ring_bcast]()
